# 计算任务
from pvsim import ship
from pvsim import wave
from pvsim import geo
from pvsim import pv
import toml
import time
import numpy as np
from scipy.integrate import RK45, dblquad
from scipy import constants
from tqdm import tqdm


class Task:
    """
    计算工况
    """

    omega = 0
    t = 0

    def __init__(self):
        "初始化静垫升船舶"
        with open("demo/cfg.toml", "r", encoding="utf-8") as file:
            self.cfg = toml.load(file)

        self.acv = ship.Ship(self.cfg["model"])
        self.acv.set_static_pressure()
        self.acv.set_static_volume()
        self.acv.set_posture(self.cfg["model"]["zb0"])

        self.static_acv = self.acv

        print("Static pressure (Pa):", self.acv.p.flatten())
        print("Static volume (m^3):", self.acv.vol.flatten())
        print("Static posture (m, rad, rad):", self.acv.x.flatten())

    def run(self):
        "运行各频率的运动响应"
        print("\nRunning...")
        self.wave = wave.Wave(self.cfg["time-domain"], 0)
        y0 = np.reshape([self.acv.x[:, 0], self.acv.v[:, 0]], [6])

        def fcn(t, y):
            "运动方程"
            ship2still = np.empty([4])
            wave2still = np.empty([4])
            p = np.empty([4, 1])
            y1 = np.zeros_like(y)
            cfg = self.cfg["model"]
            # 获取体积
            for i in range(4):
                ship2still[i] = (
                    geo.height(y[:3], self.acv.xc[i], self.acv.yc[i], cfg["zb0"])
                    * self.acv.area[i]
                )

            def wave_height(y, x):
                """
                波浪高度
                """
                return (
                    self.wave.amplitude
                    * np.tanh(0.5 * t)
                    * self.wave.wave_height(x, y, t)
                )

            w1 = dblquad(wave_height, cfg["xsm"], cfg["xsf"], 0, cfg["ysl"])
            w2 = dblquad(wave_height, cfg["xsm"], cfg["xsf"], cfg["ysr"], 0)
            w3 = dblquad(wave_height, cfg["xsb"], cfg["xsm"], 0, cfg["ysl"])
            w4 = dblquad(wave_height, cfg["xsb"], cfg["xsm"], cfg["ysr"], 0)

            wave2still = [w1[0], w2[0], w3[0], w4[0]]

            vol = ship2still - wave2still
            p[0, 0] = pv.v2p_fc(vol[0])
            p[1, 0] = pv.v2p_fc(vol[1])
            p[2, 0] = pv.v2p_bc(vol[2])
            p[3, 0] = pv.v2p_bc(vol[3])

            self.acv.p = p
            self.acv.vol[:, 0] = vol

            y1[:3] = y[-3:]
            y1[-3:] = np.reshape(
                np.linalg.solve(
                    self.acv.m,
                    self.acv.an @ p
                    - np.reshape([self.acv.m[0, 0] * constants.g, 0, 0], [3, 1]),
                ),
                [3],
            )

            return y1

        t_end = self.cfg["time-domain"]["stop_time"]
        for f in tqdm(self.cfg["time-domain"]["frequency"]):
            solution = RK45(fcn, 0, y0, t_end)
            self.wave = wave.Wave(self.cfg["time-domain"], f)
            self.acv = self.static_acv

            self.res_file = open("demo/f" + str(f) + "-res.csv", "w", encoding="utf-8")
            self.res_file.write("t,z,ry,rx,vz,vry,vrx\n")

            self.p_file = open("demo/f" + str(f) + "-p.csv", "w", encoding="utf-8")
            self.p_file.write("t,p1,p2,p3,p4\n")

            self.v_file = open("demo/f" + str(f) + "-v.csv", "w", encoding="utf-8")
            self.v_file.write("t,v1,v2,v3,v4\n")

            while solution.t < t_end:
                solution.step()
                # 在 res.csv 中输出
                self.res_file.write(
                    f"{solution.t},{solution.y[0]},{solution.y[1]},{solution.y[2]},{solution.y[3]},{solution.y[4]},{solution.y[5]}\n"
                )
                self.p_file.write(
                    f"{solution.t},{self.acv.p[0, 0]},{self.acv.p[1, 0]},{self.acv.p[2, 0]},{self.acv.p[3, 0]}\n"
                )
                self.v_file.write(
                    f"{solution.t},{self.acv.vol[0, 0]},{self.acv.vol[1, 0]},{self.acv.vol[2, 0]},{self.acv.vol[3, 0]}\n"
                )

            self.res_file.close()
            self.p_file.close()
            self.v_file.close()
